# predict_service.py
import torch
import joblib
import numpy as np
import json
from fastapi import FastAPI, Request
from pydantic import BaseModel
from models.price_predictor import PricePredictor
import uvicorn

NUMERIC_COLUMNS = ["预算", "流行度", "平均评分", "评分人数", "已上映天数", "入座率", "影院评分"]

app = FastAPI()
print(">>>> 正在运行的新文件 <<<<: predict_service.py")

def apply_adjustments(base_price, city_level=None, session=None, user_tags=None, movie_tags=None):
    adjustment = 0.0

    if city_level is not None:
        city_adjust = {1: 0.0, 2: -2, 3: -4, 4: -6, 5: -8}
        adjustment += city_adjust.get(city_level, 0.0)

    if session is not None:
        session_adjust = {"morning": -3, "noon": 0, "afternoon": 4, "evening": 3}
        adjustment += session_adjust.get(session, 0.0)

    if user_tags and movie_tags:
        overlap = len(set(user_tags) & set(movie_tags))
        adjustment += overlap * 1.0

    adjustment = max(min(adjustment, 10), -10)
    return base_price + adjustment

def predict_price_from_features(feature_dict, city_level=None, session=None,
                                user_tags=None, movie_tags=None,
                                model_path="movie-ticket-bidding/checkpoints/ResidualBlockGELU/best_model.pt",
                                config_path="movie-ticket-bidding/config/ResidualBlockGELU.json",
                                scaler_path="movie-ticket-bidding/data/processed/scaler.pkl"):
    for key in NUMERIC_COLUMNS:
        if key not in feature_dict:
            raise ValueError(f"缺少特征: {key}")

    input_vals = np.array([feature_dict[k] for k in NUMERIC_COLUMNS]).reshape(1, -1)
    scaler = joblib.load(scaler_path)
    input_scaled = scaler.transform(input_vals)
    x = torch.tensor(input_scaled, dtype=torch.float32)

    with open(config_path, "r") as f:
        model_cfg = json.load(f)["model"]

    input_dim = x.shape[1]
    model = PricePredictor(config_path, input_dim)
    model.load_state_dict(torch.load(model_path, weights_only=True))
    model.eval()

    with torch.no_grad():
        output = model(x)
        base_price = output.item()

    final_price = apply_adjustments(base_price, city_level, session, user_tags, movie_tags)
    return final_price


@app.post("/predict")
async def predict(request: Request):
    print(">>>> 正在运行的新文件 <<<<: predict_service.py")

    data = await request.json()
    try:
        feature_dict = {key: data[key] for key in NUMERIC_COLUMNS}
        city_level = data.get("city_level")
        session = data.get("session")
        user_tags = data.get("user_tags", [])
        movie_tags = data.get("movie_tags", [])

        price = predict_price_from_features(
            feature_dict=feature_dict,
            city_level=city_level,
            session=session,
            user_tags=user_tags,
            movie_tags=movie_tags
        )

        # 最低票价限制
        if price < 25:
            price = 25

        return {"predicted_price": round(price, 2)}
    except Exception as e:
        return {"error": str(e)}

if __name__ == "__main__":
    uvicorn.run("predict_service:app", host="0.0.0.0", port=8000, reload=True)
